import os
import numpy as np
import cv2
import re
from concurrent.futures import ThreadPoolExecutor


def load_flow(flow_dir, row, col):
    filename = f"tr{row}-tc{col}.npy"
    filepath = os.path.join(flow_dir, filename)
    return np.load(filepath, mmap_mode='r')


def load_image(image_dir, row, col):
    pattern = re.compile(r'(.*)tr\d+-tc\d+\.png')
    for filename in os.listdir(image_dir):
        if re.search(f'tr{row}-tc{col}', filename):
            match = pattern.match(filename)
            if match:
                prefix = match.group(1)
                filepath = os.path.join(image_dir, filename)
                return cv2.imread(filepath), prefix, filename
    return None, None, None


def save_flow(flow, row, col, output_dir):
    filename = f"tr{row}-tc{col}.npy"
    filepath = os.path.join(output_dir, filename)
    np.save(filepath, flow)
    print(f"Saved {filename} to {filepath}")


def save_image(image, filename, output_dir):
    filepath = os.path.join(output_dir, filename)
    cv2.imwrite(filepath, image)
    print(f"Saved {filename} to {filepath}")


def process_subgrid(subgrid_row, subgrid_col, rows, cols, flow_dir, image_dir, flow_output_dir, image_output_dir,
                    flow_tile_size, image_tile_size, flow_subgrid_rows, flow_subgrid_cols, row_overlap, col_overlap,
                    prefix):
    subgrid_id = (subgrid_row, subgrid_col)
    flow_combined = np.zeros((flow_subgrid_rows * flow_tile_size, flow_subgrid_cols * flow_tile_size, 2),
                             dtype=np.float32)
    image_combined = np.zeros((flow_subgrid_rows * image_tile_size, flow_subgrid_cols * image_tile_size, 3),
                              dtype=np.uint8)

    for i in range(subgrid_row, subgrid_row + flow_subgrid_rows):
        for j in range(subgrid_col, subgrid_col + flow_subgrid_cols):
            flow_part = load_flow(flow_dir, i, j)
            row_offset = (i - subgrid_row) * flow_tile_size
            col_offset = (j - subgrid_col) * flow_tile_size
            flow_combined[row_offset:row_offset + flow_tile_size, col_offset:col_offset + flow_tile_size, :] = flow_part

            image_part, temp_prefix, filename = load_image(image_dir, i, j)
            if prefix is None:
                prefix = temp_prefix
            image_combined[row_offset:row_offset + image_tile_size, col_offset:col_offset + image_tile_size,
            :] = image_part

    save_flow(flow_combined, (subgrid_row - 1) // (flow_subgrid_rows - row_overlap) + 1,
              (subgrid_col - 1) // (flow_subgrid_cols - col_overlap) + 1, flow_output_dir)
    save_image(image_combined,
               f"{prefix}tr{(subgrid_row - 1) // (flow_subgrid_rows - row_overlap) + 1}-tc{(subgrid_col - 1) // (flow_subgrid_cols - col_overlap) + 1}.png",
               image_output_dir)


def combine_flows_and_images(flow_dir, image_dir, flow_output_dir, image_output_dir):
    rows, cols = 37, 33
    flow_tile_size = 4096
    image_tile_size = 4096
    flow_subgrid_rows, flow_subgrid_cols = 7, 7
    row_overlap, col_overlap = 2, 2

    os.makedirs(flow_output_dir, exist_ok=True)
    os.makedirs(image_output_dir, exist_ok=True)

    prefix = None
    processed_subgrids = set()
    subgrid_tasks = []

    with ThreadPoolExecutor(max_workers=8) as executor:
        for subgrid_row in range(1, rows + 1, flow_subgrid_rows - row_overlap):
            if subgrid_row + flow_subgrid_rows - 1 > rows:
                subgrid_row = rows - flow_subgrid_rows + 1
            for subgrid_col in range(1, cols + 1, flow_subgrid_cols - col_overlap):
                if subgrid_col + flow_subgrid_cols - 1 > cols:
                    subgrid_col = cols - flow_subgrid_cols + 1

                subgrid_id = (subgrid_row, subgrid_col)
                if subgrid_id in processed_subgrids:
                    continue
                processed_subgrids.add(subgrid_id)

                task = executor.submit(process_subgrid, subgrid_row, subgrid_col, rows, cols, flow_dir, image_dir,
                                       flow_output_dir, image_output_dir, flow_tile_size, image_tile_size,
                                       flow_subgrid_rows, flow_subgrid_cols, row_overlap, col_overlap, prefix)
                subgrid_tasks.append(task)

        for task in subgrid_tasks:
            task.result()


# 示例用法
flow_dir = '/root/data/deep_align/32nm_flow_fine_align_to_8nm(3)'  # 替换为实际光流场文件夹路径
image_dir = '/root/data/deep_align/11866'  # 替换为实际图片文件夹路径
flow_output_dir = '/root/data/deep_align/11867_flow_big_combine_8nm(3)'  # 替换为实际光流场合并输出文件夹路径
image_output_dir = '/root/data/deep_align/deep_align/11866_big_combine_8nm(3)'  # 替换为实际图片合并输出文件夹路径

combine_flows_and_images(flow_dir, image_dir, flow_output_dir, image_output_dir)
